package ink.mint.convolution;

import ink.mint.activation.ActivationType;
import ink.mint.matrix.MatrixExtensions;

import java.io.Serializable;
import java.util.ArrayList;
import java.util.List;

import static java.lang.Math.floorDiv;

public class ConvolutionLayer implements Serializable {
    ConvolutionNeural convolutionNeural;
    List<ConvolutionNeural> convolutionNeuralList;
    int convolutionNeuralNumber;
    MatrixExtensions matrix;

    public ConvolutionLayer() {
        convolutionNeuralList = new ArrayList<>();
    }

    public void setConvolutionNeuralNumber(int convolutionNeuralNumber) {
        this.convolutionNeuralNumber = convolutionNeuralNumber;
        matrix = new MatrixExtensions();
        createConvolutionLayer();
    }

    public ConvolutionLayer input(double[][][] input) {
        int ir = input.length;
        int k = convolutionNeuralNumber / ir;
        int count = convolutionNeuralNumber % ir;
        int runK = count == 0 ? k : k + 1;
        int inputI = 0;

        if (ir == 1) {
            for (int i = 0; i < runK; i++) {
                convolutionNeural = convolutionNeuralList.get(i);
                convolutionNeural.input(input[0]);
            }
        } else {
            for (int j = 0; j < convolutionNeuralNumber; j++) {
                if (inputI >= ir) {
                    inputI = 0;
                }
                convolutionNeural = convolutionNeuralList.get(j);
                convolutionNeural.input(input[inputI]);
                inputI++;
            }
        }
        return this;
    }

    private ConvolutionLayer createConvolutionLayer() {
        for (int i = 0; i < convolutionNeuralNumber; i++) {
            convolutionNeuralList.add(new ConvolutionNeural());
        }
        return this;
    }

    public ConvolutionLayer setConvolutionLayerParameter(int kernelSize, int kernelStride, int convolutionPadding, ActivationType activationType) {
        for (int i = 0; i < convolutionNeuralNumber; i++) {
            convolutionNeural = convolutionNeuralList.get(i);
            convolutionNeural.setConvolutionParameter(kernelSize, kernelStride, convolutionPadding, activationType);
        }
        return this;
    }

    public ConvolutionLayer setPoolingLayerParameter(int poolingSize, int poolingStride, int poolingPadding, PoolingType poolingType) {
        for (int i = 0; i < convolutionNeuralNumber; i++) {
            convolutionNeural = convolutionNeuralList.get(i);
            convolutionNeural.setPoolingParameter(poolingSize, poolingStride, poolingPadding, poolingType);
        }
        return this;
    }

    public ConvolutionLayer convolution(double learningRate,boolean isTrain) {
        double[][] kernelDelta;
        double[][] kernel = new double[0][];
        double[][] newKernel = new double[0][];
        double[][] learningKernel = new double[0][];
        double bias = 0;
        double biasDelta;
        double newBias = 0;
        for (int i = 0; i < convolutionNeuralNumber; i++) {
            convolutionNeural = convolutionNeuralList.get(i);
            if (convolutionNeural.getKernelDelta() != null) {
                kernel = convolutionNeural.getKernel();
                kernelDelta = convolutionNeural.getKernelDelta();
                learningKernel = matrix.multi(kernelDelta, learningRate);
                newKernel = matrix.sub(kernel, learningKernel);
                bias = convolutionNeural.getBias();
                biasDelta = convolutionNeural.getBiasDelta();
                newBias = bias - (biasDelta * learningRate);
            }
            if(isTrain){
                convolutionNeural.convolution(newKernel, newBias,isTrain);
            }else {
                convolutionNeural.convolution(kernel, bias,isTrain);
            }

        }
        return this;
    }

    public ConvolutionLayer activate() {
        for (int i = 0; i < convolutionNeuralNumber; i++) {
            convolutionNeural = convolutionNeuralList.get(i);
            convolutionNeural.activate();
        }
        return this;
    }

    public ConvolutionLayer pooling() {
        for (int i = 0; i < convolutionNeuralNumber; i++) {
            convolutionNeural = convolutionNeuralList.get(i);
            convolutionNeural.pooling();
        }
        return this;
    }

    public double[][][] ConvolutionLayerOutput() {
        double[][] poolingMap = convolutionNeuralList.get(0).getPoolingMap();
        int pr = poolingMap.length;
        int pc = poolingMap[0].length;
        double[][][] result = new double[convolutionNeuralNumber][pr][pc];
        for (int i = 0; i < convolutionNeuralNumber; i++) {
            convolutionNeural = convolutionNeuralList.get(i);
            poolingMap = convolutionNeural.getPoolingMap();
            result[i] = poolingMap;
        }
        return result;
    }

    public ConvolutionNeural getConvolutionNeural() {
        return convolutionNeural;
    }

    public void setConvolutionNeural(ConvolutionNeural convolutionNeural) {
        this.convolutionNeural = convolutionNeural;
    }

    public List<ConvolutionNeural> getConvolutionNeuralList() {
        return convolutionNeuralList;
    }

    public void setConvolutionNeuralList(List<ConvolutionNeural> convolutionNeuralList) {
        this.convolutionNeuralList = convolutionNeuralList;
    }

    public int getConvolutionNeuralNumber() {
        return convolutionNeuralNumber;
    }
}
